import time
from ultralytics import YOLO
import cv2
 
# 加载预训练的YOLOv8n模型
# model = YOLO('./yolov8n.pt')
model = YOLO('./buwanjia2.pt')
 
# 在图片列表上运行批处理推理
results = model(['./img_naima_55.jpg'], stream=True)  # 返回 Results 对象生成器

# 初始化一个空列表来存储推理结果
tuili_jieguo = []
# 遍历推理结果生成器
for result in results:
    img = result.plot()
    cv2.imshow('Predicted Image', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    if len(result.boxes.cls) > 0:  # 如果检测到了物体
        # 遍历每个检测到的物体
        for i in range(len(result.boxes.cls)):
            # 获取类别ID
            leibie_id = int(result.boxes.cls[i].item())
            # 获取类别的名称
            leibie = result.names[leibie_id]

            # 获取物体的相似度（置信度）
            xiangsidu = result.boxes.conf[i].item()

            # 获取物体的边界框坐标值，返回的是左上角和右下角坐标
            zuobiao = result.boxes.xyxy[i].tolist()
            print(leibie)
            # 存入列表中
            tuili_jieguo.append({
                '类别': leibie,
                '相似度': xiangsidu,
                '坐标': zuobiao
            })

# 打印结果
if len(tuili_jieguo) > 0:
    for info in tuili_jieguo:
        print(info)
else:
    print('没有找到任何元素')